import torch.nn

import config
import dataset


def init_weights(m):
    classname = m.__class__.__name__
    if hasattr(m, 'weight') and classname.find('Conv') != -1:

        torch.nn.init.xavier_normal_(m.weight)

    elif classname.find('BatchNorm2d') != -1:
        torch.nn.init.normal_(m.weight, 1.0, 0.02)
        torch.nn.init.constant_(m.bias, 0.0)


def get_dataset(args: config.DeepLearningYaml, transform):
    if args.dataset == 'ade20k':
        train = dataset.ade20k.Ade20kSeg(root=args.dir, images_path='training',
                                         seg_path='training',
                                         image_size=(args.image_size, args.image_size), transform=transform)
        val = dataset.ade20k.Ade20kSeg(root=args.dir, images_path='validation',
                                       seg_path='validation',
                                       image_size=(args.image_size, args.image_size), transform=transform)
    elif args.dataset == 'voc2012':
        train = dataset.voc2012.get_voc_2012(args.dir, 'train')
        val = dataset.voc2012.get_voc_2012(args.dir, 'val')
    else:
        raise Exception('No validate dataset name')
    return train, val
